import os

import socket
import socks
from langchain import SQLDatabase, OpenAI, SQLDatabaseChain
from langchain.prompts.prompt import PromptTemplate

try:  # 离开ide也能运行
    from conf.config import config, logger, BASE_DIR, database_uri
    from utils.constants import _DEFAULT_TEMPLATE
except ImportError:
    import os
    import sys
    sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))  # 离开IDE也能正常导入自己定义的包
    from conf.config import config, logger, BASE_DIR, database_uri
    from utils.constants import _DEFAULT_TEMPLATE


PROMPT = PromptTemplate(
    input_variables=["input", "table_info", "dialect", "top_k"],
    template=_DEFAULT_TEMPLATE,
)


def query(question: str):

    os.environ['OPENAI_API_KEY'] = config.get("open_ai_api_key")
    if config.get("proxy_enable"):  # 被墙
        socks.set_default_proxy(socks.SOCKS5, config.get("proxy_host"), config.get("proxy_port_socket"))
        socket.socket = socks.socksocket

    db = SQLDatabase.from_uri(database_uri)

    llm = OpenAI(temperature=0, max_tokens=2048)

    db_chain = SQLDatabaseChain(
        llm=llm,
        database=db,
        prompt=PROMPT,
        top_k=config.get("query_house_price_top_k"),
        verbose=True,
    )
    answer = db_chain.run(question).strip()
    return answer


def main():
    question = "2023年3月份深圳市所有区的新房和二手房的房价，用表格格式返回，包含年月"
    logger.info(f"问题：{question}")
    answer = query(question)
    logger.info(f"答复：{answer}")


if __name__ == '__main__':
    main()
